#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
改进的Neo4j查询工具
功能：基于文件名匹配的智能查询
"""

from neo4j import GraphDatabase
import jieba
import re
import sys


class ImprovedNeo4jQuery:
    # def __init__(self, uri="bolt://localhost:7687", user="neo4j", password="changsha2025"):
    def __init__(self, uri="bolt://192.168.1.169:7687", user="neo4j", password="Pa@ssw0rd"):
        try:
            self.driver = GraphDatabase.driver(uri, auth=(user, password))
            print(f"✅ 成功连接到Neo4j数据库")
        except Exception as e:
            print(f"❌ 连接Neo4j失败: {e}")
            self.driver = None
    
    def close(self):
        if self.driver:
            self.driver.close()
    
    def show_all_files(self):
        """显示所有文件"""
        if not self.driver:
            return []
        
        with self.driver.session() as session:
            cypher = "MATCH (n:File) RETURN n.name as filename LIMIT 25"
            print(f"📝 执行查询: {cypher}")
            
            try:
                result = session.run(cypher)
                files = [record['filename'] for record in result]
                
                print(f"\n📁 找到 {len(files)} 个文件:")
                for i, filename in enumerate(files, 1):
                    print(f"  {i}. {filename}")
                
                return files
            except Exception as e:
                print(f"❌ 查询失败: {e}")
                return []
    
    def query_by_product_code(self, product_code):
        """根据产品编号查询完整决策链路"""
        if not self.driver:
            return []

        # 构建文件名
        filename = f"{product_code}.txt"

        with self.driver.session() as session:
            results = []

            # 查询1: 获取完整的决策链路
            cypher1 = """
                MATCH (f:File {name: $filename})-[:CONTAINS]->(l:LineRecord)
                RETURN l.content as full_content, l.line_number as line_number
                ORDER BY l.line_number
            """
            rendered_cypher1 = cypher1.replace('$filename', f'"{filename}"')
            print(f"📝 查询1 - 完整决策链路:")
            print(f"   {rendered_cypher1}")

            try:
                result = session.run(cypher1, filename=filename)
                for record in result:
                    results.append({
                        'type': 'decision_chain',
                        'content': record['full_content'],
                        'line_number': record['line_number'],
                        'source': 'file_match'
                    })
            except Exception as e:
                print(f"❌ 决策链路查询失败: {e}")

            # 查询2: 获取决策节点的层次结构
            cypher2 = """
                MATCH (f:File {name: $filename})-[:CONTAINS]->(l:LineRecord)-[:CONTAINS_NODE]->(n:DecisionNode)
                RETURN l.line_number as line_number,
                       collect({content: n.content, type: n.type, level: n.level}) as nodes
                ORDER BY l.line_number
            """
            rendered_cypher2 = cypher2.replace('$filename', f'"{filename}"')
            print(f"📝 查询2 - 决策节点层次:")
            print(f"   {rendered_cypher2}")

            try:
                result = session.run(cypher2, filename=filename)
                for record in result:
                    results.append({
                        'type': 'decision_structure',
                        'line_number': record['line_number'],
                        'nodes': record['nodes'],
                        'source': 'structure_match'
                    })
            except Exception as e:
                print(f"❌ 决策结构查询失败: {e}")

            return results
    
    def query_by_keyword(self, keyword):
        """根据关键词查询完整链路"""
        if not self.driver:
            return []

        with self.driver.session() as session:
            results = []

            # 查询1: 查找包含关键词的完整决策链路
            cypher1 = """
                MATCH (l:LineRecord)
                WHERE l.content CONTAINS $keyword
                RETURN l.content as full_content, l.filename as filename, l.line_number as line_number
                ORDER BY l.filename, l.line_number
                LIMIT 10
            """
            rendered_cypher1 = cypher1.replace('$keyword', f'"{keyword}"')
            print(f"📝 查询1 - 关键词完整链路搜索:")
            print(f"   {rendered_cypher1}")

            try:
                result = session.run(cypher1, keyword=keyword)
                for record in result:
                    results.append({
                        'type': 'keyword_chain',
                        'content': record['full_content'],
                        'filename': record['filename'],
                        'line_number': record['line_number'],
                        'matched_keyword': keyword,
                        'source': 'keyword_chain'
                    })
            except Exception as e:
                print(f"❌ 关键词链路查询失败: {e}")

            # 查询2: 查找包含关键词的决策节点及其上下文
            cypher2 = """
                MATCH (n:DecisionNode)
                WHERE n.content CONTAINS $keyword
                MATCH (l:LineRecord)-[:CONTAINS_NODE]->(n)
                RETURN l.content as full_content, n.content as node_content, n.type as node_type,
                       l.filename as filename, l.line_number as line_number
                ORDER BY l.filename, l.line_number
                LIMIT 10
            """
            rendered_cypher2 = cypher2.replace('$keyword', f'"{keyword}"')
            print(f"📝 查询2 - 关键词节点上下文:")
            print(f"   {rendered_cypher2}")

            try:
                result = session.run(cypher2, keyword=keyword)
                for record in result:
                    results.append({
                        'type': 'keyword_context',
                        'full_content': record['full_content'],
                        'node_content': record['node_content'],
                        'node_type': record['node_type'],
                        'filename': record['filename'],
                        'line_number': record['line_number'],
                        'matched_keyword': keyword,
                        'source': 'keyword_context'
                    })
            except Exception as e:
                print(f"❌ 关键词上下文查询失败: {e}")

            return results
    
    def query_by_filename_only(self, filename):
        """只查询指定文件下的数据"""
        if not self.driver:
            return []

        # 确保文件名以.txt结尾
        if not filename.endswith('.txt'):
            filename = f"{filename}.txt"

        with self.driver.session() as session:
            results = []

            # 查询该文件的所有完整决策链路
            cypher = """
                MATCH (f:File {name: $filename})-[:CONTAINS]->(l:LineRecord)
                RETURN l.content as full_content, l.line_number as line_number
                ORDER BY l.line_number
            """
            rendered_cypher = cypher.replace('$filename', f'"{filename}"')
            print(f"📝 文件专用查询:")
            print(f"   {rendered_cypher}")

            try:
                result = session.run(cypher, filename=filename)
                for record in result:
                    results.append({
                        'type': 'file_only',
                        'content': record['full_content'],
                        'line_number': record['line_number'],
                        'filename': filename,
                        'source': 'file_exclusive'
                    })
                print(f"✅ 找到 {len(results)} 条记录")
            except Exception as e:
                print(f"❌ 文件查询失败: {e}")

            return results

    def query_by_keyword_with_file_filter(self, keyword, filename=None):
        """根据关键词查询，可选择限制在特定文件内"""
        if not self.driver:
            return []

        with self.driver.session() as session:
            results = []

            # 构建文件过滤条件
            file_filter = ""
            if filename:
                file_filter = f' AND l.filename = "{filename}"'

            # 查询1: 查找包含关键词的完整决策链路
            cypher1 = f"""
                MATCH (l:LineRecord)
                WHERE l.content CONTAINS $keyword{file_filter}
                RETURN l.content as full_content, l.filename as filename, l.line_number as line_number
                ORDER BY l.filename, l.line_number
                LIMIT 10
            """
            rendered_cypher1 = cypher1.replace('$keyword', f'"{keyword}"')
            print(f"📝 查询1 - 关键词完整链路搜索:")
            print(f"   {rendered_cypher1}")

            try:
                result = session.run(cypher1, keyword=keyword)
                for record in result:
                    results.append({
                        'type': 'keyword_chain',
                        'content': record['full_content'],
                        'filename': record['filename'],
                        'line_number': record['line_number'],
                        'matched_keyword': keyword,
                        'source': 'keyword_chain'
                    })
            except Exception as e:
                print(f"❌ 关键词链路查询失败: {e}")

            # 查询2: 查找包含关键词的决策节点及其上下文
            cypher2 = f"""
                MATCH (n:DecisionNode)
                WHERE n.content CONTAINS $keyword{file_filter.replace('l.filename', 'n.filename') if filename else ''}
                MATCH (l:LineRecord)-[:CONTAINS_NODE]->(n)
                RETURN l.content as full_content, n.content as node_content, n.type as node_type,
                       l.filename as filename, l.line_number as line_number
                ORDER BY l.filename, l.line_number
                LIMIT 10
            """
            rendered_cypher2 = cypher2.replace('$keyword', f'"{keyword}"')
            print(f"📝 查询2 - 关键词节点上下文:")
            print(f"   {rendered_cypher2}")

            try:
                result = session.run(cypher2, keyword=keyword)
                for record in result:
                    results.append({
                        'type': 'keyword_context',
                        'full_content': record['full_content'],
                        'node_content': record['node_content'],
                        'node_type': record['node_type'],
                        'filename': record['filename'],
                        'line_number': record['line_number'],
                        'matched_keyword': keyword,
                        'source': 'keyword_context'
                    })
            except Exception as e:
                print(f"❌ 关键词上下文查询失败: {e}")

            return results

    def smart_query(self, query_text):
        """智能查询"""
        print(f"🔍 智能查询: '{query_text}'")
        print("=" * 60)

        # 检查是否是纯文件名查询
        query_upper = query_text.upper().strip()

        # 情况1: 直接输入文件名 (如 "TD504.txt" 或 "TD504")
        if (query_upper.endswith('.TXT') or
            re.match(r'^[A-Z]{1,3}\d{3,4}$', query_upper)):

            filename = query_upper
            if not filename.endswith('.txt'):
                filename = f"{filename}.txt"

            print(f"🎯 检测到文件名专用查询: {filename}")
            print(f"📁 只查询该文件下的数据，所有查询都限制在此文件内")

            results = self.query_by_filename_only(filename)
            return results

        # 情况2: 复合查询 (包含产品编号和关键词)
        # 提取产品编号
        product_codes = re.findall(r'([A-Z]{1,3}\d{3,4})', query_text.upper())

        # 提取关键词
        words = list(jieba.cut(query_text))
        keywords = [w.strip() for w in words if len(w.strip()) > 1 and w.strip() not in ['的', '是', '有', '和', '怎么', '什么']]

        print(f"🎯 识别的产品编号: {product_codes}")
        print(f"🔍 提取的关键词: {keywords}")

        all_results = []

        # 如果有产品编号，确定文件范围
        target_filename = None
        if product_codes:
            target_filename = f"{product_codes[0]}.txt"
            print(f"📁 限制查询范围到文件: {target_filename}")

        # 策略1: 如果有产品编号，只按产品编号查询文件内容，不再进行额外查询
        if product_codes:
            product_code = product_codes[0]  # 只处理第一个产品编号
            print(f"\n📦 查询产品文件: {product_code}")
            results = self.query_by_product_code(product_code)
            all_results.extend(results)

            # 如果有其他关键词，只在该文件内搜索
            non_product_keywords = [k for k in keywords if k.upper() not in product_codes]
            if non_product_keywords:
                print(f"\n🔍 在文件 {target_filename} 内搜索关键词: {non_product_keywords}")
                for keyword in non_product_keywords[:3]:
                    if len(keyword) > 1:
                        results = self.query_by_keyword_with_file_filter(keyword, target_filename)
                        all_results.extend(results)
        else:
            # 策略2: 没有产品编号时，按关键词进行跨文件查询
            print(f"\n🔍 跨文件搜索关键词:")
            for keyword in keywords[:3]:  # 限制关键词数量
                if len(keyword) > 1:
                    print(f"   关键词: {keyword}")
                    results = self.query_by_keyword(keyword)
                    all_results.extend(results)

        # 去重
        unique_results = self._deduplicate_results(all_results)

        return unique_results
    
    def _deduplicate_results(self, results):
        """去重结果"""
        seen = set()
        unique_results = []
        
        for result in results:
            # 创建唯一标识
            key = f"{result.get('filename', '')}-{result.get('line_number', '')}-{result.get('content', '')}"
            
            if key not in seen:
                seen.add(key)
                unique_results.append(result)
        
        return unique_results
    
    def display_results(self, results, query):
        """显示查询结果 - 重点显示完整决策链路"""
        if not results:
            print(f"\n❌ 没有找到关于 '{query}' 的相关信息")
            return

        print(f"\n🎯 查询结果 (共{len(results)}条):")
        print("=" * 80)

        # 检查是否是文件专用查询
        file_only_results = [r for r in results if r.get('type') == 'file_only']

        if file_only_results:
            # 文件专用查询结果
            print(f"\n📁 文件专用查询结果 - {file_only_results[0].get('filename', 'N/A')} ({len(file_only_results)}条):")
            for i, result in enumerate(file_only_results, 1):
                print(f"\n{i}. 行{result.get('line_number', 'N/A')}:")
                print(f"   {result['content']}")
            return

        # 按来源分组显示
        decision_chains = [r for r in results if r.get('type') == 'decision_chain']
        decision_structures = [r for r in results if r.get('type') == 'decision_structure']
        keyword_chains = [r for r in results if r.get('type') == 'keyword_chain']
        keyword_contexts = [r for r in results if r.get('type') == 'keyword_context']

        # 显示完整决策链路
        if decision_chains:
            print(f"\n🔗 完整决策链路 ({len(decision_chains)}条):")
            for i, result in enumerate(decision_chains[:10], 1):
                print(f"\n{i}. 行{result.get('line_number', 'N/A')}:")
                print(f"   完整链路: {result['content']}")

        # 显示决策结构
        if decision_structures:
            print(f"\n🏗️  决策结构分析 ({len(decision_structures)}条):")
            for i, result in enumerate(decision_structures[:5], 1):
                print(f"\n{i}. 行{result.get('line_number', 'N/A')} 决策节点:")
                nodes = result.get('nodes', [])
                for j, node in enumerate(nodes, 1):
                    print(f"   {j}. [{node.get('type', 'N/A')}] 级别{node.get('level', 'N/A')}: {node.get('content', 'N/A')}")

        # 显示关键词完整链路
        if keyword_chains:
            print(f"\n🔍 关键词匹配的完整链路 ({len(keyword_chains)}条):")
            for i, result in enumerate(keyword_chains[:10], 1):
                print(f"\n{i}. 文件: {result.get('filename', 'N/A')} 行{result.get('line_number', 'N/A')}")
                print(f"   匹配词: {result.get('matched_keyword', 'N/A')}")

                # 高亮关键词
                content = result['content']
                keyword = result.get('matched_keyword', '')
                if keyword:
                    content = content.replace(keyword, f"【{keyword}】")
                print(f"   完整链路: {content}")

        # 显示关键词上下文
        if keyword_contexts:
            print(f"\n📋 关键词节点上下文 ({len(keyword_contexts)}条):")
            for i, result in enumerate(keyword_contexts[:5], 1):
                print(f"\n{i}. 文件: {result.get('filename', 'N/A')} 行{result.get('line_number', 'N/A')}")
                print(f"   匹配词: {result.get('matched_keyword', 'N/A')}")
                print(f"   节点类型: {result.get('node_type', 'N/A')}")

                # 高亮关键词
                node_content = result.get('node_content', '')
                keyword = result.get('matched_keyword', '')
                if keyword:
                    node_content = node_content.replace(keyword, f"【{keyword}】")
                print(f"   匹配节点: {node_content}")
                print(f"   完整链路: {result.get('full_content', 'N/A')}")


def main():
    """主函数"""
    if len(sys.argv) > 1:
        # 命令行模式
        query = ' '.join(sys.argv[1:])
        
        query_tool = ImprovedNeo4jQuery()
        if not query_tool.driver:
            return
        
        try:
            if query.lower() == 'files':
                # 显示所有文件
                query_tool.show_all_files()
            else:
                # 智能查询
                results = query_tool.smart_query(query)
                query_tool.display_results(results, query)
        finally:
            query_tool.close()
    
    else:
        # 交互模式
        print("🔍 改进的Neo4j查询工具")
        print("=" * 40)
        print("💡 输入 'files' 查看所有文件")
        print("💡 输入产品编号如 'TD504' 查看产品信息")
        print("💡 输入关键词进行搜索")
        print("💡 输入 'quit' 退出")
        print("=" * 40)
        
        query_tool = ImprovedNeo4jQuery()
        if not query_tool.driver:
            return
        
        try:
            while True:
                query = input("\n🔍 请输入查询: ").strip()
                
                if query.lower() in ['quit', 'exit', '退出']:
                    print("👋 再见!")
                    break
                
                if not query:
                    continue
                
                if query.lower() == 'files':
                    query_tool.show_all_files()
                else:
                    results = query_tool.smart_query(query)
                    query_tool.display_results(results, query)
        
        finally:
            query_tool.close()



if __name__ == "__main__":
    main()
